import numpy as np
import pandas as pd
import logging
from enum import Enum
from typing import Dict, Any, Optional, List, Tuple

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('PositionSizer')

class SizingMethod(Enum):
    """仓位计算方法枚举"""
    FIXED_SIZE = 1            # 固定大小
    PERCENTAGE_OF_EQUITY = 2  # 基于总资产百分比
    RISK_PER_TRADE = 3        # 基于每笔交易风险
    VOLATILITY_BASED = 4      # 基于波动率
    OPTIMAL_F = 5             # 最优F值法
    KELLY_CRITERION = 6       # 凯利准则
    MARGIN_BASED = 7          # 基于保证金

class PositionSizer:
    """动态仓位算法模块，负责计算最优仓位大小"""
    def __init__(self, config=None):
        self.config = config or {
            'default_sizing_method': SizingMethod.RISK_PER_TRADE,
            'fixed_size': 100,                    # 固定大小法的默认值
            'equity_percentage': 0.02,            # 资产百分比法的默认值（2%）
            'risk_per_trade': 0.02,               # 每笔交易风险的默认值（2%）
            'max_position_percent': 0.1,          # 单个品种最大仓位比例（10%）
            'max_total_exposure': 0.5,            # 总风险敞口上限（50%）
            'volatility_target': 0.15,            # 目标波动率（15%年化）
            'risk_free_rate': 0.03,               # 无风险利率
            'minimum_position_size': 10,          # 最小仓位大小
            'maximum_position_size': 100000,      # 最大仓位大小
            'portfolio_margin_rate': 0.1,         # 组合保证金比例（10%）
            'instrument_margin_rates': {},        # 各品种保证金比例
            'debug_mode': False
        }
        
        # 账户信息
        self.account_info = {
            'equity': 0.0,
            'available_funds': 0.0,
            'margin_used': 0.0,
            'open_positions': {},
            'total_exposure': 0.0
        }
        
        # 历史仓位记录
        self.position_history = []
        
        # 性能统计
        self.performance_stats = {
            'total_positions_calculated': 0,
            'avg_position_size': 0,
            'max_position_size': 0,
            'min_position_size': float('inf'),
            'position_size_variability': 0
        }
    
    def update_account_info(self, account_info: Dict[str, Any]):
        """更新账户信息"""
        # 更新账户信息
        self.account_info.update(account_info)
        
        # 计算总风险敞口（如果没有提供）
        if 'total_exposure' not in account_info and self.account_info.get('open_positions'):
            self.account_info['total_exposure'] = self._calculate_total_exposure()
        
        logger.debug(f"Account info updated: Equity=${self.account_info['equity']:.2f}")
    
    def _calculate_total_exposure(self) -> float:
        """计算总风险敞口"""
        open_positions = self.account_info.get('open_positions', {})
        total_value = 0.0
        
        for symbol, position in open_positions.items():
            qty = position.get('qty', 0)
            price = position.get('price', 0)
            total_value += abs(qty) * price
        
        equity = self.account_info.get('equity', 1.0)  # 避免除以0
        return total_value / equity if equity > 0 else 0
    
    def calculate_position_size(self, symbol: str, signal_info: Dict[str, Any], method: Optional[SizingMethod] = None) -> Dict[str, Any]:
        """计算仓位大小"""
        try:
            # 如果没有指定方法，使用默认方法
            if method is None:
                method = self.config['default_sizing_method']
            
            # 记录计算开始时间
            self.performance_stats['total_positions_calculated'] += 1
            
            # 根据不同方法计算仓位大小
            if method == SizingMethod.FIXED_SIZE:
                position_size = self._calculate_fixed_size()
            elif method == SizingMethod.PERCENTAGE_OF_EQUITY:
                position_size = self._calculate_percentage_of_equity(symbol, signal_info)
            elif method == SizingMethod.RISK_PER_TRADE:
                position_size = self._calculate_risk_per_trade(symbol, signal_info)
            elif method == SizingMethod.VOLATILITY_BASED:
                position_size = self._calculate_volatility_based(symbol, signal_info)
            elif method == SizingMethod.OPTIMAL_F:
                position_size = self._calculate_optimal_f(symbol, signal_info)
            elif method == SizingMethod.KELLY_CRITERION:
                position_size = self._calculate_kelly_criterion(symbol, signal_info)
            elif method == SizingMethod.MARGIN_BASED:
                position_size = self._calculate_margin_based(symbol, signal_info)
            else:
                raise ValueError(f"Unsupported sizing method: {method}")
            
            # 应用仓位限制
            position_size = self._apply_position_limits(symbol, position_size)
            
            # 记录历史
            self._record_position_history(symbol, method, position_size, signal_info)
            
            # 更新性能统计
            self._update_performance_stats(position_size)
            
            return {
                'success': True,
                'symbol': symbol,
                'position_size': position_size,
                'method': method.name,
                'signal_info': signal_info,
                'timestamp': pd.Timestamp.now().timestamp()
            }
            
        except Exception as e:
            logger.error(f"Error calculating position size for {symbol}: {str(e)}")
            return {
                'success': False,
                'error': str(e),
                'symbol': symbol,
                'timestamp': pd.Timestamp.now().timestamp()
            }
    
    def _calculate_fixed_size(self) -> int:
        """固定大小法"""
        return self.config['fixed_size']
    
    def _calculate_percentage_of_equity(self, symbol: str, signal_info: Dict[str, Any]) -> int:
        """基于总资产百分比法"""
        # 获取账户权益
        equity = self.account_info.get('equity', 0.0)
        
        if equity <= 0:
            return self.config['minimum_position_size']
        
        # 获取当前价格
        current_price = signal_info.get('current_price', 0.0)
        
        if current_price <= 0:
            return self.config['minimum_position_size']
        
        # 计算基于权益百分比的仓位价值
        position_value = equity * self.config['equity_percentage']
        
        # 计算股数
        position_size = int(position_value / current_price)
        
        return position_size
    
    def _calculate_risk_per_trade(self, symbol: str, signal_info: Dict[str, Any]) -> int:
        """基于每笔交易风险法"""
        # 获取账户权益
        equity = self.account_info.get('equity', 0.0)
        
        if equity <= 0:
            return self.config['minimum_position_size']
        
        # 获取当前价格和止损价格
        current_price = signal_info.get('current_price', 0.0)
        stop_loss_price = signal_info.get('stop_loss_price', 0.0)
        
        if current_price <= 0:
            return self.config['minimum_position_size']
        
        # 如果没有提供止损价格，使用ATR计算
        if stop_loss_price <= 0:
            atr = signal_info.get('atr', 0.0)
            if atr > 0:
                # 基于ATR设置止损（通常为1-3倍ATR）
                atr_multiplier = signal_info.get('atr_multiplier', 2.0)
                stop_loss_distance = atr * atr_multiplier
            else:
                # 如果没有ATR数据，使用默认百分比止损
                default_stop_percent = signal_info.get('default_stop_percent', 0.02)
                stop_loss_distance = current_price * default_stop_percent
            
            # 根据交易方向设置止损价格
            direction = signal_info.get('direction', 'BUY')
            if direction == 'BUY':
                stop_loss_price = current_price - stop_loss_distance
            else:
                stop_loss_price = current_price + stop_loss_distance
        
        # 计算每单位风险
        risk_per_unit = abs(current_price - stop_loss_price)
        
        if risk_per_unit <= 0:
            return self.config['minimum_position_size']
        
        # 计算总风险金额
        total_risk_amount = equity * self.config['risk_per_trade']
        
        # 计算仓位大小
        position_size = int(total_risk_amount / risk_per_unit)
        
        return position_size
    
    def _calculate_volatility_based(self, symbol: str, signal_info: Dict[str, Any]) -> int:
        """基于波动率法"""
        # 获取账户权益
        equity = self.account_info.get('equity', 0.0)
        
        if equity <= 0:
            return self.config['minimum_position_size']
        
        # 获取当前价格和波动率
        current_price = signal_info.get('current_price', 0.0)
        volatility = signal_info.get('volatility', 0.0)
        
        if current_price <= 0 or volatility <= 0:
            return self.config['minimum_position_size']
        
        # 目标波动率（年化）
        target_volatility = self.config['volatility_target']
        
        # 计算基于波动率的仓位大小
        # 公式: 仓位大小 = (目标波动率 / 当前波动率) * (账户权益 / 当前价格)
        position_size = int((target_volatility / volatility) * (equity / current_price))
        
        return position_size
    
    def _calculate_optimal_f(self, symbol: str, signal_info: Dict[str, Any]) -> int:
        """最优F值法"""
        # 获取账户权益
        equity = self.account_info.get('equity', 0.0)
        
        if equity <= 0:
            return self.config['minimum_position_size']
        
        # 获取当前价格
        current_price = signal_info.get('current_price', 0.0)
        
        if current_price <= 0:
            return self.config['minimum_position_size']
        
        # 获取历史交易结果（盈亏比）
        # 实际应用中，这里应该使用该策略或品种的历史交易数据
        win_rate = signal_info.get('win_rate', 0.5)  # 胜率
        risk_reward_ratio = signal_info.get('risk_reward_ratio', 1.5)  # 盈亏比
        
        # 计算最优F值
        # 公式: f = (胜率 * 盈亏比 - (1 - 胜率)) / 盈亏比
        optimal_f = (win_rate * risk_reward_ratio - (1 - win_rate)) / risk_reward_ratio
        optimal_f = max(0.0, min(optimal_f, 1.0))  # 限制在0-1之间
        
        # 计算仓位大小
        # 公式: 仓位大小 = (最优F值 * 账户权益) / (最大单笔亏损 * 当前价格)
        max_loss_per_unit = signal_info.get('max_loss_per_unit', current_price * 0.02)  # 每单位最大亏损
        position_size = int((optimal_f * equity) / (max_loss_per_unit * current_price))
        
        return position_size
    
    def _calculate_kelly_criterion(self, symbol: str, signal_info: Dict[str, Any]) -> int:
        """凯利准则法"""
        # 获取账户权益
        equity = self.account_info.get('equity', 0.0)
        
        if equity <= 0:
            return self.config['minimum_position_size']
        
        # 获取当前价格
        current_price = signal_info.get('current_price', 0.0)
        
        if current_price <= 0:
            return self.config['minimum_position_size']
        
        # 获取历史交易结果
        win_rate = signal_info.get('win_rate', 0.5)  # 胜率
        avg_win = signal_info.get('avg_win', current_price * 0.03)  # 平均盈利
        avg_loss = signal_info.get('avg_loss', current_price * 0.02)  # 平均亏损
        
        if avg_loss <= 0:
            return self.config['minimum_position_size']
        
        # 计算赔率
        odds = avg_win / avg_loss
        
        # 计算凯利百分比
        # 公式: Kelly% = (胜率 * 赔率 - (1 - 胜率)) / 赔率
        kelly_percentage = (win_rate * odds - (1 - win_rate)) / odds
        kelly_percentage = max(0.0, min(kelly_percentage, 1.0))  # 限制在0-1之间
        
        # 为了安全起见，通常使用半凯利或四分之一凯利
        kelly_fraction = signal_info.get('kelly_fraction', 0.5)  # 半凯利
        adjusted_kelly_percentage = kelly_percentage * kelly_fraction
        
        # 计算仓位大小
        position_value = equity * adjusted_kelly_percentage
        position_size = int(position_value / current_price)
        
        return position_size
    
    def _calculate_margin_based(self, symbol: str, signal_info: Dict[str, Any]) -> int:
        """基于保证金法"""
        # 获取可用资金
        available_funds = self.account_info.get('available_funds', 0.0)
        
        if available_funds <= 0:
            return self.config['minimum_position_size']
        
        # 获取当前价格
        current_price = signal_info.get('current_price', 0.0)
        
        if current_price <= 0:
            return self.config['minimum_position_size']
        
        # 获取该品种的保证金比例
        margin_rate = self.config['instrument_margin_rates'].get(symbol, self.config['portfolio_margin_rate'])
        
        # 计算可用于该交易的资金（考虑多元化）
        diversification_factor = signal_info.get('diversification_factor', 0.5)  # 多元化因子
        usable_funds = available_funds * diversification_factor
        
        # 计算最大可交易价值
        max_trade_value = usable_funds / margin_rate
        
        # 计算仓位大小
        position_size = int(max_trade_value / current_price)
        
        return position_size
    
    def _apply_position_limits(self, symbol: str, position_size: int) -> int:
        """应用仓位限制"""
        # 获取账户权益
        equity = self.account_info.get('equity', 0.0)
        current_price = self._get_current_price(symbol)
        
        # 应用最小和最大仓位限制
        position_size = max(self.config['minimum_position_size'], min(self.config['maximum_position_size'], position_size))
        
        # 应用单个品种最大仓位比例限制
        if equity > 0 and current_price > 0:
            max_position_value = equity * self.config['max_position_percent']
            max_position_size = int(max_position_value / current_price)
            position_size = min(position_size, max_position_size)
        
        # 应用总风险敞口限制
        if equity > 0 and current_price > 0:
            current_exposure = self.account_info.get('total_exposure', 0.0)
            new_position_value = position_size * current_price
            new_exposure = (current_exposure * equity + new_position_value) / equity
            
            if new_exposure > self.config['max_total_exposure']:
                # 计算可添加的最大仓位以不超过总风险敞口限制
                max_additional_value = equity * (self.config['max_total_exposure'] - current_exposure)
                max_additional_size = int(max_additional_value / current_price)
                position_size = min(position_size, max_additional_size)
        
        # 确保仓位大小为正数
        position_size = max(0, position_size)
        
        return position_size
    
    def _get_current_price(self, symbol: str) -> float:
        """获取当前价格"""
        # 实际应用中，这里应该从市场数据服务获取最新价格
        # 这里返回一个模拟价格
        return 150.0  # 假设当前价格为150
    
    def _record_position_history(self, symbol: str, method: SizingMethod, position_size: int, signal_info: Dict[str, Any]):
        """记录仓位历史"""
        history_record = {
            'timestamp': pd.Timestamp.now().timestamp(),
            'symbol': symbol,
            'method': method.name,
            'position_size': position_size,
            'signal_info': signal_info,
            'account_equity': self.account_info.get('equity', 0.0),
            'total_exposure': self.account_info.get('total_exposure', 0.0)
        }
        
        self.position_history.append(history_record)
        
        # 限制历史记录大小
        max_history_size = 1000
        if len(self.position_history) > max_history_size:
            self.position_history.pop(0)
    
    def _update_performance_stats(self, position_size: int):
        """更新性能统计"""
        # 更新平均仓位大小（使用指数加权移动平均）
        self.performance_stats['avg_position_size'] = (
            self.performance_stats['avg_position_size'] * 0.9 + position_size * 0.1
        )
        
        # 更新最大和最小仓位大小
        self.performance_stats['max_position_size'] = max(self.performance_stats['max_position_size'], position_size)
        self.performance_stats['min_position_size'] = min(self.performance_stats['min_position_size'], position_size)
        
        # 计算仓位大小变异性（简单标准差）
        if len(self.position_history) >= 10:
            recent_sizes = [record['position_size'] for record in self.position_history[-10:]]
            self.performance_stats['position_size_variability'] = np.std(recent_sizes)
    
    def get_position_history(self, symbol: str = None, start_time: float = None, end_time: float = None) -> List[Dict[str, Any]]:
        """获取仓位历史记录"""
        filtered_history = self.position_history
        
        # 按股票代码过滤
        if symbol:
            filtered_history = [record for record in filtered_history if record['symbol'] == symbol]
        
        # 按时间范围过滤
        if start_time:
            filtered_history = [record for record in filtered_history if record['timestamp'] >= start_time]
        
        if end_time:
            filtered_history = [record for record in filtered_history if record['timestamp'] <= end_time]
        
        return filtered_history
    
    def get_performance_stats(self) -> Dict[str, Any]:
        """获取性能统计信息"""
        return self.performance_stats.copy()
    
    def calculate_portfolio_allocation(self, signals: Dict[str, Dict[str, Any]], method: Optional[SizingMethod] = None) -> Dict[str, Any]:
        """计算投资组合分配"""
        allocation = {}
        total_weight = 0.0
        
        try:
            # 计算每个品种的初始权重
            for symbol, signal_info in signals.items():
                # 计算仓位大小
                result = self.calculate_position_size(symbol, signal_info, method)
                
                if result.get('success', False):
                    position_size = result['position_size']
                    current_price = signal_info.get('current_price', 0.0)
                    
                    if current_price > 0:
                        position_value = position_size * current_price
                        allocation[symbol] = {
                            'position_size': position_size,
                            'position_value': position_value,
                            'weight': 0.0,  # 初始化为0，后面再计算
                            'signal_info': signal_info
                        }
                        total_weight += position_value
            
            # 计算权重
            if total_weight > 0:
                for symbol, alloc in allocation.items():
                    alloc['weight'] = alloc['position_value'] / total_weight
            
            return {
                'success': True,
                'allocation': allocation,
                'total_weight': total_weight,
                'timestamp': pd.Timestamp.now().timestamp()
            }
            
        except Exception as e:
            logger.error(f"Error calculating portfolio allocation: {str(e)}")
            return {
                'success': False,
                'error': str(e),
                'timestamp': pd.Timestamp.now().timestamp()
            }
    
    def optimize_position_sizing(self, symbol: str, backtest_results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """优化仓位大小参数"""
        try:
            # 实际应用中，这里应该使用回测结果来优化仓位大小参数
            # 这里简单返回当前参数
            return {
                'success': True,
                'symbol': symbol,
                'optimal_params': {
                    'risk_per_trade': self.config['risk_per_trade'],
                    'volatility_target': self.config['volatility_target'],
                    'kelly_fraction': 0.5
                },
                'timestamp': pd.Timestamp.now().timestamp()
            }
            
        except Exception as e:
            logger.error(f"Error optimizing position sizing for {symbol}: {str(e)}")
            return {
                'success': False,
                'error': str(e),
                'symbol': symbol,
                'timestamp': pd.Timestamp.now().timestamp()
            }
    
    def set_sizing_parameters(self, **kwargs):
        """设置仓位计算参数"""
        for key, value in kwargs.items():
            if key in self.config:
                self.config[key] = value
        
        logger.info(f"Updated sizing parameters: {kwargs}")
        return self.config

# 示例使用
if __name__ == "__main__":
    # 创建PositionSizer实例
    sizer = PositionSizer()
    
    # 更新账户信息
    sizer.update_account_info({
        'equity': 100000.0,
        'available_funds': 80000.0,
        'margin_used': 20000.0
    })
    
    # 准备信号信息
    signal_info = {
        'current_price': 150.0,
        'stop_loss_price': 145.0,
        'direction': 'BUY',
        'volatility': 0.2,
        'win_rate': 0.6,
        'risk_reward_ratio': 2.0
    }
    
    # 使用不同方法计算仓位大小
    risk_per_trade_result = sizer.calculate_position_size('AAPL', signal_info, SizingMethod.RISK_PER_TRADE)
    print(f"Risk per trade position size: {risk_per_trade_result}")
    
    volatility_based_result = sizer.calculate_position_size('AAPL', signal_info, SizingMethod.VOLATILITY_BASED)
    print(f"Volatility based position size: {volatility_based_result}")
    
    kelly_result = sizer.calculate_position_size('AAPL', signal_info, SizingMethod.KELLY_CRITERION)
    print(f"Kelly criterion position size: {kelly_result}")
    
    # 获取性能统计
    stats = sizer.get_performance_stats()
    print(f"Performance stats: {stats}")
    
    # 计算投资组合分配
    portfolio_signals = {
        'AAPL': signal_info,
        'MSFT': {
            'current_price': 300.0,
            'stop_loss_price': 290.0,
            'direction': 'BUY',
            'volatility': 0.15,
            'win_rate': 0.55,
            'risk_reward_ratio': 1.8
        },
        'GOOG': {
            'current_price': 2800.0,
            'stop_loss_price': 2700.0,
            'direction': 'BUY',
            'volatility': 0.18,
            'win_rate': 0.58,
            'risk_reward_ratio': 1.9
        }
    }
    
    allocation_result = sizer.calculate_portfolio_allocation(portfolio_signals)
    print(f"Portfolio allocation: {allocation_result}")